import re
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from collections import OrderedDict

def _bn_function_factory(relu, conv):
    def bn_function(*inputs):
        concated_features = torch.cat(inputs, 1)
        bottleneck_output = conv(relu(concated_features))
        return bottleneck_output

    return bn_function

class _DenseLayer(nn.Sequential):
    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate, memory_efficient=False):
        super(_DenseLayer, self).__init__()
        #self.add_module('norm1', nn.BatchNorm1d(num_input_features)),
        self.add_module('relu1', nn.ReLU(inplace=True)),
        self.add_module('conv1', nn.Conv1d(num_input_features, bn_size *
                                           growth_rate, kernel_size=1, stride=1,
                                           bias=False)),
        #self.add_module('norm2', nn.BatchNorm1d(bn_size * growth_rate)),
        self.add_module('relu2', nn.ReLU(inplace=True)),
        self.add_module('conv2', nn.Conv1d(bn_size * growth_rate, growth_rate,
                                           kernel_size=3, stride=1, padding=1,
                                           bias=False)),
        self.drop_rate = drop_rate
        self.memory_efficient = memory_efficient

    def forward(self, *prev_features):
        bn_function = _bn_function_factory(self.relu1, self.conv1)
        if self.memory_efficient and any(prev_feature.requires_grad for prev_feature in prev_features):
            bottleneck_output = cp.checkpoint(bn_function, *prev_features)
        else:
            bottleneck_output = bn_function(*prev_features)
        new_features = self.conv2(self.relu2(bottleneck_output))
        if self.drop_rate > 0:
            new_features = F.dropout(new_features, p=self.drop_rate,
                                     training=self.training)
        return new_features


class _DenseBlock(nn.Module):
    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate, memory_efficient=False):
        super(_DenseBlock, self).__init__()
        for i in range(num_layers):
            layer = _DenseLayer(
                num_input_features + i * growth_rate,
                growth_rate=growth_rate,
                bn_size=bn_size,
                drop_rate=drop_rate,
                memory_efficient=memory_efficient,
            )
            self.add_module('denselayer%d' % (i + 1), layer)

    def forward(self, init_features):
        features = [init_features]
        for name, layer in self.named_children():
            new_features = layer(*features)
            features.append(new_features)
        return torch.cat(features, 1)




def _load_state_dict(model, model_url, progress):
    # '.'s are no longer allowed in module names, but previous _DenseLayer
    # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
    # They are also in the checkpoints in model_urls. This pattern is used
    # to find such keys.
    pattern = re.compile(
        r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')

    state_dict = load_state_dict_from_url(model_url, progress=progress)
    for key in list(state_dict.keys()):
        res = pattern.match(key)
        if res:
            new_key = res.group(1) + res.group(2)
            state_dict[new_key] = state_dict[key]
            del state_dict[key]
    model.load_state_dict(state_dict)


def _densenet(arch, growth_rate, block_config, num_init_features, pretrained, progress,
              **kwargs):
    model = DenseNet(growth_rate, block_config, num_init_features, **kwargs)
    if pretrained:
        _load_state_dict(model, model_urls[arch], progress)
    return model

class TickDenseNet_slim(nn.Module):
    def __init__(self, input_feature_size=68, target_size=2,
                 num_init_features=64, block_config=[6, 12]):
        super(TickDenseNet_slim, self).__init__()
        
        # First convolution
        self.features = nn.Sequential(OrderedDict([
            ('conv0', nn.Conv1d(input_feature_size, num_init_features, kernel_size=7, stride=2,
                                padding=3, bias=False)),
           # ('norm0', nn.BatchNorm1d(num_init_features)),
            ('relu0', nn.ReLU(inplace=True)),
            ('pool0', nn.MaxPool1d(kernel_size=3, stride=2, padding=1)),
        ]))
        
        num_features = num_init_features
        
        for i, num_layers in enumerate(block_config):
            block = _DenseBlock(
                num_layers=num_layers,
                num_input_features=num_features,
                bn_size=4,
                growth_rate=32,
                drop_rate=0,
            )
            self.features.add_module('denseblock{}'.format(i + 1), block)
            num_features = num_features + num_layers * 32

        print(num_features)
        # final stage 
        #self.features.add_module('norm5', nn.BatchNorm1d(num_features))
        # lstm layer 
        self.lstm_layer = nn.LSTM(
            input_size=num_features, 
            hidden_size = 256, 
            num_layers=1,
            dropout=0.2,
            bidirectional = False
        )
        
        # self.aux = nn.Linear(num_features, target_size)
        self.aux = nn.Linear(256, target_size)
        
    def forward(self, input_tensor):
        out = self.features(input_tensor)

        out, _ = self.lstm_layer(out.permute((2, 0, 1)))
        out = out.permute((1, 2, 0))

        # out = F.relu(out, inplace=True)
        out = F.adaptive_avg_pool1d(out, (1)).squeeze()
        

        out = self.aux(out)
        
        #out, (out_h, out_c) = self.lstm_layer(out.permute((2, 0, 1)))
        


        #out = self.aux(out_c.squeeze())
        return out
    
if __name__ == "__main__":
    import numpy as np
    input_tensor = torch.from_numpy(np.random.random((10, 200, 68))).float().permute((0, 2, 1))

    model = TickDenseNet_slim()

    # print(model)

    cc = model(input_tensor)






